#!/usr/bin/env python
# encoding: utf-8
"""
@author: youfeng
@email: youfeng243@163.com
@license: Apache Licence
@file: extract_company.py
@time: 2018/1/9 15:59
"""

import sys

sys.path.append('..')
sys.path.append('../..')
from common import util
from common.mongo import MongDb
from logger import Logger

MONGO_DB_TARGET = {
    'host': '172.16.215.2',
    'port': 40042,
    'db': 'company_data',
    'username': 'work',
    'password': 'haizhi'
}

MONGO_DB_SOURCE = {
    'host': '172.16.215.16',
    'port': 40042,
    'db': 'app_data',
    'username': 'work',
    'password': 'haizhi'
}

# 日志模块
log = Logger('extract_company.log').get_logger()

# 工商库
source_db = MongDb(MONGO_DB_SOURCE['host'], MONGO_DB_SOURCE['port'], MONGO_DB_SOURCE['db'],
                   MONGO_DB_SOURCE['username'],
                   MONGO_DB_SOURCE['password'], log=log)

# 种子库
target_db = MongDb(MONGO_DB_TARGET['host'], MONGO_DB_TARGET['port'], MONGO_DB_TARGET['db'],
                   MONGO_DB_TARGET['username'],
                   MONGO_DB_TARGET['password'], log=log)


def judge_type(company):
    feature_list = ['银行', '公司', '集团', '实验室', '研究院', '工程', '有限', '责任', '国有', '股份', '全民', '上市', '独资', '控股', '投资',
                    '所有制',
                    '企业']
    for feature in feature_list:
        if feature in company:
            return '企业'

    return '个体户'


def get_capital_grade(registered_capital):
    if registered_capital <= 5000000:
        return 0
    if registered_capital > 5000000 and 5000000 <= 10000000:
        return 1
    if registered_capital > 10000000 and registered_capital <= 20000000:
        return 2
    if registered_capital > 20000000 and 20000000 <= 50000000:
        return 3
    return 4


def find(item, count_dict):
    #  注册资本评判 500以下 500 - 1000  1000 - 2000 2000 - 5000 5000以上
    # 913507817264517242 [4:6]
    SINGLE_MAX = 30
    company = item.get('company')
    registered_capital = item.get('registered_capital')
    business_status = item.get('business_status')
    unified_social_credit_code = item.get('unified_social_credit_code')
    if company is None:
        return None

    if registered_capital is None:
        return None

    try:
        registered_capital = float(registered_capital)
    except Exception as e:
        return None

    if business_status is None:
        return None

    if business_status == '':
        return None

    if unified_social_credit_code is None:
        return None

    if len(unified_social_credit_code) != 18:
        return None

    city = unified_social_credit_code[4:6]
    company_type = judge_type(company)
    capital_grade = get_capital_grade(registered_capital)

    data_item = {
        'city': city,
        'type': company_type,
        'registered_capital': registered_capital,
        'capital_grade': capital_grade,
        'status': business_status
    }
    # 城市
    # 企业类型
    # 注册资本等级
    # 企业状态
    if city in count_dict:
        if company_type in count_dict[city]:
            if capital_grade in count_dict[city][company_type]:
                if business_status in count_dict[city][company_type][capital_grade]:
                    if count_dict[city][company_type][capital_grade][business_status] >= SINGLE_MAX:
                        return None

                    count_dict[city][company_type][capital_grade][business_status] += 1
                    return data_item

                count_dict[city][company_type][capital_grade][business_status] = 1
                return data_item

            count_dict[city][company_type][capital_grade] = {}
            count_dict[city][company_type][capital_grade][business_status] = 1
            return data_item

        count_dict[city][company_type] = {}
        count_dict[city][company_type][capital_grade] = {}
        count_dict[city][company_type][capital_grade][business_status] = 1
        return data_item

    count_dict[city] = {}
    count_dict[city][company_type] = {}
    count_dict[city][company_type][capital_grade] = {}
    count_dict[city][company_type][capital_grade][business_status] = 1
    return data_item


# {
#     # 城市
#     'city': '',
#     # 企业类型
#     'type': '',
#     # 注册资本
#     'registered_capital': '',
#     # 企业状态
#     'status': '',
# }

def main():
    source_table_name = 'enterprise_data_gov'
    target_table_name = 'new_gansu_company_list_20180109'

    point = 0
    total = 0
    result_list = []
    count_dict = {}
    for item in source_db.traverse_batch_field(source_table_name, {},
                                               ['company',
                                                'province',
                                                'registered_capital',
                                                'business_status', 'unified_social_credit_code']):
        point += 1
        if point % 1000 == 0:
            log.info("当前扫描进度: {} find: {}".format(point, total))

        province = item.get('province')
        if province is None:
            continue

        if not (province == '甘肃' or province == '甘肃省'):
            continue

        if total >= 10000:
            break

        # {
        #     # 城市
        #     'city': '',
        #     # 企业类型
        #     'type': '',
        #     # 注册资本
        #     'registered_capital': '',
        #     # 企业状态
        #     'status': '',
        # }
        data_item = find(item, count_dict)
        if data_item is not None:
            data_item['_id'] = item.get('company')
            data_item['_in_time'] = util.get_now_time()
            data_item['crawl_status'] = 0
            total += 1
            result_list.append(data_item)
            if len(result_list) >= 500:
                target_db.insert_batch_data(target_table_name, result_list)
                del result_list[:]

    target_db.insert_batch_data(target_table_name, result_list)
    del result_list[:]


if __name__ == '__main__':
    main()
